# just for test
import sys
sys.path.append("E:\work\python_work")

import math
import random
from pyrobot.bean.map.point import Point
from pyrobot.bean.map.mapstructure import MapStructure  
from pyrobot.bean.map.mapstructure import Cell  
#import copy

class Particle:
    def __init__(self):
        self.v = Point()
        self.pos = Point()
        self.best = Point()
        self.fitness = 0

"""
This class works on the geography map, the distance unit is MM. 
"""
class PSO:
    def __init__(self, particleNum, iterNum, robot, globalMap, freeEdgesList, sampleDist):
        self.particleNum = particleNum
        self.iterNum = iterNum      
        self.robot = robot
        self.globalMap = globalMap
        self.sampleDist = sampleDist
        
        # data format transformation
        self.freeEdgesList = []
        print "freeEdgesList ", freeEdgesList
        print "*****************************"
        for freeedge in freeEdgesList:
#            print freeedge
#            print freeedge[0][0]
#            print freeedge[0][1]
#            print freeedge[1][0]
#            print freeedge[1][1]
            
            cellX, cellY = self.globalMap.geoCord2CellCord(freeedge[0][0], freeedge[0][1])
            cell1 = Cell(cellX, cellY, "obstacle")
            cellX, cellY = self.globalMap.geoCord2CellCord(freeedge[1][0], freeedge[1][1])
            cell2 = Cell(cellX, cellY, "obstacle")
            
            freeEdgeSubSet = []
            freeEdgeSubSet.append(cell1)
            freeEdgeSubSet.append(cell2)
            self.freeEdgesList.append(freeEdgeSubSet)
        
        # sonar ray parameters
#        self.rayMaxDist = robot.__dict__['range'].getMaxvalue()
        # for test
        self.rayMaxDist = 8000 
        
#        self.raysSet = copy.deepcopy(robot.__dict__['range'])
#        self.raysSet.units = 'METERS'
#        self.raysSet = robot.__dict__['range']
        # initialize the particles
        self.__initParticles__()    
        # keep track of the best particle
        self.gbest = Particle()
        # termination condition
        self.terminateDist = 500
        # learning rate
        self.c1 = 1.0
        self.c2 = 1.0
     
    def __randomPos__(self):
        # pick one randomly from the free edge list        
        
#        print "len(self.freeEdgesList) ", len(self.freeEdgesList)
        
        edge = random.randint(0, len(self.freeEdgesList)-1)        
#        print "len(self.freeEdgesList) ", len(self.freeEdgesList)
#        print "edge ", edge       
        obstacle = self.freeEdgesList[edge]
        # get the cell coordinate
        point1 = self.globalMap.getCellCoordinate(obstacle[0])
        point2 = self.globalMap.getCellCoordinate(obstacle[1])
        midpoint = ((point1[0]+point2[0])/2, (point1[1]+point2[1])/2)
        midpointX = midpoint[0]
        midpointY = midpoint[1]

        #
        x = random.uniform(-1, 1) * self.sampleDist + midpointX
        y = random.uniform(-1, 1) * self.sampleDist + midpointY 
        
        print "midpointX ", midpointX
        print "midpointY ", midpointY
        print "self.globalMap.getGeoCoordinateValue(midpointX, midpointY) ", self.globalMap.getGeoCoordinateValue(midpointX, midpointY)
        
        count = 0
        while self.globalMap.getGeoCoordinateValue(x, y) == "obstacle" or \
              self.globalMap.getGeoCoordinateValue(x, y) == 0 :
#        while self.globalMap.getGeoCoordinateValue(x, y) == "obstacle":
            x = random.uniform(-1, 1) * self.sampleDist + midpointX
            y = random.uniform(-1, 1) * self.sampleDist + midpointY 
            count += 1
            if count > 10:
                x = midpointX + 10
                y = midpointY + 10
                print "break"
                break
            print "repeat generating sample"
#            print "x ", x
#            print "y ", y
            
        dir = math.atan(float(y-midpointY)/float(x-midpointX))   
        return x, y, dir
    
    def __randomPos__(self, obstacle):
        # get the middle point of the free edge
        point1 = self.globalMap.getCellCoordinate(obstacle[0])
        point2 = self.globalMap.getCellCoordinate(obstacle[1])
        midpoint = ((point1[0]+point2[0])/2, (point1[1]+point2[1])/2)
        midpointX = midpoint[0]
        midpointY = midpoint[1]
        midpoint = Point(midpointX, midpointY)
#        midpoint = self.__adjustPosition__(midpoint)

        # get the sample area
        x = y = 0.0
        if self.globalMap.getGeoCoordinateValue(midpoint.x - self.sampleDist, midpoint.y) == 0:            
            x = random.uniform(1, 2) * self.sampleDist + midpoint.x
            if point1[0] > point2[0]:
                if point1[1] > point2[1]:
                    y = random.uniform(-1.5, -0.5) * self.sampleDist + midpoint.y
                else:
                    y = random.uniform(0.5, 1.5) * self.sampleDist + midpoint.y
            else:
                if point1[1] > point2[1]:
                    y = random.uniform(-1.5, -0.5) * self.sampleDist + midpoint.y
                else:
                    y = random.uniform(0.5, 1.5) * self.sampleDist + midpoint.y
        else:
            x = random.uniform(-2, -1) * self.sampleDist + midpoint.x
            if point1[0] > point2[0]:
                if point1[1] > point2[1]:
                    y = random.uniform(-1.5, -0.5) * self.sampleDist + midpoint.y
                else:
                    y = random.uniform(0.5, 1.5) * self.sampleDist + midpoint.y
            else:
                if point1[1] > point2[1]:
                    y = random.uniform(-1.5, -0.5) * self.sampleDist + midpoint.y
                else:
                    y = random.uniform(0.5, 1.5) * self.sampleDist + midpoint.y
            
        dir = math.atan(float(y-midpointY)/float(x-midpointX))   
        return x, y, dir
    
    def __conflictFreeOrNot__(self, particle):
        # test 8 directions
        testPoint = Point()
        
        x = particle.pos.x + 400
        y = particle.pos.y + 400
#        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle" \
#         or self.globalMap.getGeoCoordinateValue(x, y) == 0:
        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle":
            return False
        x = particle.pos.x + 400
        y = particle.pos.y - 400
#        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle" \
#         or self.globalMap.getGeoCoordinateValue(x, y) == 0:
        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle":
            return False
        x = particle.pos.x - 400
        y = particle.pos.y - 400
#        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle" \
#         or self.globalMap.getGeoCoordinateValue(x, y) == 0:
        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle":
            return False
        x = particle.pos.x - 400
        y = particle.pos.y + 400
#        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle" \
#         or self.globalMap.getGeoCoordinateValue(x, y) == 0:
        if self.globalMap.getGeoCoordinateValue(x, y) == "obstacle":
            return False
        
        return True
    
    def __initParticles__(self):
        self.particles = []        
        for i in range(len(self.freeEdgesList)):
            for j in range(self.particleNum):
                p = Particle()
                # generate pos
                p.pos.x, p.pos.y, p.pos.dir = self.__randomPos__(self.freeEdgesList[i])                
#                while not self.__conflictFreeOrNot__(p):
#                    print "conflict, need to be re-sample"
#                    p.pos.x, p.pos.y, p.pos.dir = self.__randomPos__(self.freeEdgesList[i])

                self.particles.append(p) 

    def __restrictPosition__(self, pos):
        
#        print "pos.x ", pos.x
#        print "pos.y ", pos.y        
        # if the position is out of range
        minX, maxX, minY, maxY = self.globalMap.getCurrentXYRange() 
        if pos.x > maxX:
            pos.x = maxX
        elif pos.x < minX:
            pos.x = minX
        if pos.y > maxY:
            pos.y = maxY
        elif pos.y < minY:
            pos.y = minY   
             
        # search for free space   
        offset = 0
        while self.globalMap.getGeoCoordinateValue(pos.x, pos.y) == 0 or \
              self.globalMap.getGeoCoordinateValue(pos.x, pos.y) == "obstacle":
#        while self.globalMap.getGeoCoordinateValue(pos.x, pos.y) == "obstacle":
            offset += 50                      
            if self.globalMap.getGeoCoordinateValue(pos.x + offset, pos.y) == "free" :
                return Point(pos.x + offset, pos.y)
            elif self.globalMap.getGeoCoordinateValue(pos.x - offset, pos.y) == "free":
                return Point(pos.x - offset, pos.y)
            elif self.globalMap.getGeoCoordinateValue(pos.x, pos.y + offset) == "free":
                return Point(pos.x, pos.y + offset)
            elif self.globalMap.getGeoCoordinateValue(pos.x, pos.y - offset) == "free":
                return Point(pos.x, pos.y - offset)
        # otherwise
        return pos
    
    def __adjustPosition__(self, pos):
        
#        print "pos.x ", pos.x
#        print "pos.y ", pos.y        
        # if the position is out of range
        minX, maxX, minY, maxY = self.globalMap.getCurrentXYRange() 
        if pos.x > maxX:
            pos.x = maxX
        elif pos.x < minX:
            pos.x = minX
        if pos.y > maxY:
            pos.y = maxY
        elif pos.y < minY:
            pos.y = minY   
             
        # search for free space   
        offset = 0
        while self.globalMap.getGeoCoordinateValue(pos.x, pos.y) == 0 or \
              self.globalMap.getGeoCoordinateValue(pos.x, pos.y) == "obstacle":
            offset += 50
#            print "offset ", offset
#            print "self.globalMap.getGeoCoordinateValue(pos.x, pos.y) ", self.globalMap.getGeoCoordinateValue(pos.x, pos.y)
            if self.globalMap.getGeoCoordinateValue(pos.x + offset, pos.y) == "free":
                return Point(pos.x + offset, pos.y)
            elif self.globalMap.getGeoCoordinateValue(pos.x - offset, pos.y) == "free":
                return Point(pos.x - offset, pos.y)
            elif self.globalMap.getGeoCoordinateValue(pos.x, pos.y + offset) == "free":
                return Point(pos.x, pos.y + offset)
            elif self.globalMap.getGeoCoordinateValue(pos.x, pos.y - offset) == "free":
                return Point(pos.x, pos.y - offset)
        # otherwise
        return pos
    
    def __terminateOrNot__(self):
        for particle in self.particles:
#            print "self.gbest.pos.dist(particle.pos) ", self.gbest.pos.dist(particle.pos)
            if self.gbest.pos.dist(particle.pos) > self.terminateDist:
                return False
        return True
    
    """
    
    """
    def __getFitnessTest__(self, pos):
        dist = math.sqrt( math.pow((pos.x-self.robot.x), 2) + math.pow((pos.y-self.robot.y), 2))
        
#        print "dist ", dist
        
        feature1 = math.exp(float(10/dist))
        
        thetaSet = [90*math.pi/180, 80*math.pi/180, 70*math.pi/180, 60*math.pi/180, 
                    50*math.pi/180, 40*math.pi/180, 30*math.pi/180, 20*math.pi/180,
                    10*math.pi/180, 0*math.pi/180, -10*math.pi/180, -20*math.pi/180,
                    -30*math.pi/180, -40*math.pi/180, -50*math.pi/180, -60*math.pi/180,
                    -70*math.pi/180, -80*math.pi/180, -90*math.pi/180]
        
        count = 0
        for i in range(thetaSet.__len__()):
            # in MM:
            theta = thetaSet[i]
            # get start point and end point
            endX = math.cos(theta+pos.dir) * self.rayMaxDist + pos.x
            endY = math.sin(theta+pos.dir) * self.rayMaxDist + pos.y
            # check if the line hits obstacles
            if not self.globalMap.hitObstacleOrNot(pos.x, pos.y, endX, endY):
                count += 1
        # set feature2
        feature2 = count
        
#        print "feature1 ", feature1
#        print "feature2 ", feature2
        
        fitness = feature1 + feature2
        return fitness
    
    
    """
    
    """
    def __getFitness__(self, pos):
        dist = math.sqrt( math.pow((pos.x-self.robot.x), 2) + math.pow((pos.y-self.robot.y), 2))
        feature1 = math.exp(float(10/dist))
        
        item = 'range'
        originalUnits = self.robot.__dict__[item].units
        self.robot.__dict__[item].units = 'METERS'
        count = 0
        for i in range(self.robot.__dict__[item].count):
            # in MM:
            offx, offy, z, theta, arc = self.robot.__dict__[item][i].geometry
            # get start point and end point
            endX = math.cos(theta+pos.dir) * self.rayMaxDist + pos.x
            endY = math.sin(theta+pos.dir) * self.rayMaxDist + pos.y
            # check if the line hits obstacles
            if not self.globalMap.hitObstacleOrNot(pos.x, pos.y, endX, endY):
                count += 1
        # restore the unit
        self.robot.__dict__[item].units = originalUnits
        # set feature2
        feature2 = count
        
        fitness = feature1 + feature2
        return fitness
    
    def getBestParticle(self):
        i = 0
        bestChanged = False
        while i < self.iterNum :
            for p in self.particles:
                # calculate value
                fitness = self.__getFitness__(p.pos)
#                fitness = self.__getFitnessTest__(p.pos)
                
#                print "fitness ", fitness
#                print "p.fitness ", p.fitness
                
                if fitness > p.fitness:
                    p.fitness = fitness
                    p.best = p.pos
         
                if fitness > self.gbest.fitness:
                    self.gbest = p
                    bestChanged = True
                # update the speed
                p1 = (p.best - p.pos)
                p2 = (self.gbest.pos - p.pos)     
#                print "p.best ", p.best
#                print "p.pos ", p.pos
#                print "self.gbest.pos ", self.gbest.pos     
#                print "p1 ", p1
#                print "p2 ", p2
                
                v = Point()          
                v.x = p.v.x + self.c1 * random.random() * p1.x \
                            + self.c2 * random.random() * p2.x
                v.y = p.v.y + self.c1 * random.random() * p1.y \
                            + self.c2 * random.random() * p2.y
                
#                v = p.v + self.c1 * random.random() * p1 + self.c2 * random.random() * p2
#                v = p.v + self.c1 * random.random() * (p.best - p.pos) \
#                        + self.c2 * random.random() * (self.gbest.pos - p.pos)
                # update the position
#                print "v ", v
#                print "before ", p.pos
                p.pos = p.pos + v 
                p.pos = self.__restrictPosition__(p.pos)
#                p.pos = self.__adjustPosition__(p.pos)
                    
#                print "after ", p.pos
#                print "-----------------------------------------------------"
            
            if self.__terminateOrNot__():
                print "terminate!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! at ", i
                break; 
#            print "+++++++++++++++++ end one iteration ++++++++++++++++++++"
            i += 1
            print "pso i ", i
            
        print "self.gbest.pos ", self.gbest.pos
            
        self.gbest.pos = self.__restrictPosition__(self.gbest.pos)
        return self.gbest.pos
    
    def getNextBestView(self, freeEdgesList):
        minDist = 100000
        robotPos = Point(self.robot.x*1000, self.robot.y*1000)
        bestPoint = Point()
        for freeedge in freeEdgesList:
            midx = ( freeedge[0][0] + freeedge[1][0] )/2
            midy = ( freeedge[0][1] + freeedge[1][1] )/2
            midPoint = Point(midx, midy)
            
            dist = midPoint.dist(robotPos)
            if dist < minDist:
                minDist = dist
                bestPoint = midPoint
        print "bestPoint ", bestPoint
        return bestPoint
 
if __name__=="__main__":
     print "hello world" 
     
#     j = 0
#     while j<10:
#         print random.random()
#         j += 1
     
     testMap = MapStructure(500, 500, 50, 50)
#     cellGeoList = [(0.10000000000001941, 316.93474435699142), (61.111632370624619, 345.97816658503604), (130.35464496032699, 357.82937238821535), (219.92713438750437, 380.73454132278562), (323.4148903646103, 385.30250677196807), (509.15751280782308, 427.1355725365604), (2532.4843843783015, 1462.0506372287682), (3837.9086781316773, 1396.8346273691675), (5187.3014697414983, 914.63753261344164), (5172.9365518950935, 0.0), (2673.5182571730702, -471.38972858562545), (1758.9469438569295, -640.15443791471739), (1673.0702244045629, -965.86764062255395), (1147.3998312067031, -962.68446650679255), (782.56683648320359, -932.49848810542437), (512.79019206401279, -887.98823687847278), (329.48974477183145, -904.94856254116598), (247.06898352957847, -1400.5947114370767), (0.10000000000008449, -1379.925817452734)]
     cellGeoList = [(0.10000000000001918, 313.45326557372221), (30.213336940553035, 344.16574032589824), (61.836653190072312, 350.08996397721404), (96.456793417861334, 359.56876602247951), (125.9761701319719, 345.7996116656592), (152.88029910011048, 327.6472094513677), (216.65922702565564, 375.07435973763523), (262.31807481627987, 374.4741171900011), (331.73011538924601, 395.21220606830695), (391.80056137683613, 391.68931137683609), (506.82065838028404, 425.17471884829729), (592.45489678967522, 414.75321438240275), (753.34756397258093, 434.86551504235069), (955.32694182235025, 445.41234979047186), (1274.044136302311, 463.6642494527236), (1262.1860327621807, 338.16543960817188), (1314.3953016408771, 231.73968142506851), (1304.0441873296065, 114.07739817892362), (1280.8424174413267, 0.0), (1285.5081694910716, -112.45570675113751), (1265.6471205968658, -223.14406184654837), (1259.9273992493054, -337.5602405823991), (1250.7411348414389, -455.18265055190443), (2221.3643484303439, -1035.7752880024732), (2315.6558065586837, -1336.8645994566079), (2168.6047079544278, -1518.3851935077728), (2275.1472696259634, -1908.976926145516), (2276.0706550360869, -2275.9594050360861), (2300.447125068004, -2741.4377751542966), (2264.2457548809052, -3233.5231431161601), (2396.8011936913217, -4151.191013715702), (2008.7097315974045, -4307.4862707271413), (1542.3816793095623, -4237.3417649222138), (1189.2520956978037, -4437.9363564038149), (791.2111785952917, -4486.5784496001288), (367.23974095004922, -4196.3951655660167), (0.1000000000001065, -1739.4256911694886)]
     
     for cellGeo in cellGeoList:
        testMap.occupy(cellGeo[0], cellGeo[1])    
     
     robot = Point()
     robot.x = 0
     robot.y = 0
     pso = PSO(5, 50, robot, testMap, 500)
     
     potition = pso.getBestParticle()
    
     print "******************"
     print "pso result ", potition
     
     
     